import os
import json
import urllib.parse
import json5
from typing import List, Dict, Any, Optional
from qwen_agent.agents import Assistant
from qwen_agent.tools.base import BaseTool, register_tool
from qwen_agent.gui import WebUI
import warnings
warnings.filterwarnings("ignore")


@register_tool('my_image_gen')
class MyImageGen(BaseTool):
    """AI图像生成工具"""
    description = 'AI 绘画（图像生成）服务，输入文本描述，返回基于文本信息绘制的图像 URL。'
    parameters = [{
        'name': 'prompt',
        'type': 'string',
        'description': '期望的图像内容的详细描述',
        'required': True
    }]

    def call(self, params: str, **kwargs) -> str:
        prompt = json5.loads(params)['prompt']
        prompt = urllib.parse.quote(prompt)
        return json5.dumps(
            {'image_url': f'https://image.pollinations.ai/prompt/{prompt}'},
            ensure_ascii=False
        )


class AdaptiveRetrievalConfig:
    """
    自适应检索配置类
    根据文档数量自动选择合适的检索方式
    """
    def __init__(self, 
                 doc_threshold: int = 50,
                 use_es_retrieval: bool = False,
                 es_config: Optional[Dict] = None):
        """
        初始化配置
        :param doc_threshold: 文档数量阈值，超过此数量使用ES检索
        :param use_es_retrieval: 是否强制使用ES检索
        :param es_config: ES配置信息
        """
        self.doc_threshold = doc_threshold
        self.use_es_retrieval = use_es_retrieval
        self.es_config = es_config or {
            "host": "https://localhost",
            "port": 9200,
            "user": "elastic",
            "password": "ncUg27UP5aMSzSElVXAX",
            "index_name": "adaptive_search_docs_index"
        }
    
    def should_use_es_retrieval(self, file_count: int) -> bool:
        """
        判断是否应该使用ES检索
        :param file_count: 文档数量
        :return: 是否使用ES检索
        """
        if self.use_es_retrieval:
            return True
        return file_count > self.doc_threshold


def init_agent_service(retrieval_config: Optional[AdaptiveRetrievalConfig] = None):
    """
    初始化具备自适应检索能力的助手服务
    :param retrieval_config: 检索配置
    """
    # 默认配置
    if retrieval_config is None:
        retrieval_config = AdaptiveRetrievalConfig()
    
    # 步骤 1: LLM 配置
    llm_cfg = {
        'model': 'qwen-plus',#'qwen-max',
        'model_server': 'dashscope',
        'api_key': os.getenv('DASHSCOPE_API_KEY'),
        'generate_cfg': {
            'top_p': 0.8
        }
    }

    # 步骤 2: 获取文件列表
    file_dir = os.path.join(os.path.dirname(__file__), 'docs')
    files = []
    if os.path.exists(file_dir):
        for file in os.listdir(file_dir):
            file_path = os.path.join(file_dir, file)
            if os.path.isfile(file_path):
                files.append(file_path)
    file_count = len(files)
    print(f'知识库文件列表: {files}')
    print(f'文档数量: {file_count}')

    # 步骤 3: 根据文档数量选择检索方式
    use_es_retrieval = retrieval_config.should_use_es_retrieval(file_count)
    
    if use_es_retrieval:
        print("使用 Elasticsearch 检索 (es_retrieval)")
        # 使用 Elasticsearch 检索配置
        rag_cfg = {
            "rag_backend": "elasticsearch",
            "es": retrieval_config.es_config,
            "parser_page_size": 500
        }
        # 工具配置 - 使用my_image_gen和code_interpreter工具
        tools_cfg = ['my_image_gen', 'code_interpreter']
    else:
        print("使用默认检索 (qwen-agent retrieval)")
        # 不使用特殊RAG配置，但必须启用本地RAG功能
        rag_cfg = {
            "rag_backend": "local",  # 添加本地RAG配置
            "parser_page_size": 500
        }
        # 工具配置 - 使用my_image_gen和code_interpreter工具
        tools_cfg = ['my_image_gen', 'code_interpreter']
    
    # 添加 tavily-mcp 网络搜索工具（仅在需要时添加）
    # 修改：将网络搜索工具放在最后，作为备选方案
    print("use_es_retrieval:", use_es_retrieval)
    if use_es_retrieval:
        tools_cfg.append({
            "mcpServers": {
                "tavily-mcp": {
                    "command": "npx",
                    "args": ["-y", "tavily-mcp@0.1.4"],
                    "env": {
                        "TAVILY_API_KEY": os.getenv('TAVILY_API_KEY', "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")
                    },
                    "disabled": False,
                    "autoApprove": []
                }
            }
        })
    print(tools_cfg)
    # 步骤 4: 系统指令
    system_instruction = '''你是一个AI助手，专门用于文档检索和问答任务。
请根据用户的问题，优先利用检索工具从本地知识库中查找最相关的信息。
如果本地知识库没有相关信息，再使用 tavily_search 工具从互联网上搜索，并结合这些信息给出专业、准确的回答。
知识库中是以前做过的示例，如果有类似的需要进行逻辑上的参考。
示例：
问题: 雇主安心保的赔偿范围
答案:雇主安心保赔偿范围：
1.	死亡赔偿金
2.	伤残赔偿金
3.	医疗费用
4.	误工费用
5.	法律诉讼费用
具体投保方案和保费请以投保页面展示和计算结果为准，不同地区会有承保方案和价格差异。

当前可用的知识库包括：
- 平安商业综合责任保险
- 雇主责任险知识库
- 企业团体综合意外险
- 财产一切险知识库
- 装修保险知识库

回答时请遵循以下原则：
1. 优先使用本地知识库的信息
2. 如果使用了网络搜索信息，请明确说明
3. 回答应准确、简洁，避免冗余信息
4. 如果不确定答案，请说明不确定，不要编造信息'''


    # 步骤 5: 创建智能体实例
    bot = Assistant(
        name="保险助手",
        description="一个用于保险的助手，能够处理保险相关的问题。优先使用内部保险知识库。", 
        llm=llm_cfg,
        system_message=system_instruction,
        function_list=tools_cfg,
        files=files,
        rag_cfg=rag_cfg  # 确保在所有情况下都设置rag_cfg
    )
    return bot, use_es_retrieval


def main(doc_threshold: int = 50, force_es_retrieval: bool = False):
    """
    启动 Web 图形界面
    :param doc_threshold: 文档数量阈值
    :param force_es_retrieval: 是否强制使用ES检索
    """
    try:
        # 创建检索配置
        es_config = {
            "host": "https://localhost",
            "port": 9200,
            "user": "elastic",
            "password": "ncUg27UP5aMSzSElVXAX",
            "index_name": "adaptive_search_docs_index"
        }
        
        retrieval_config = AdaptiveRetrievalConfig(
            doc_threshold=doc_threshold,
            use_es_retrieval=force_es_retrieval,
            es_config=es_config
        )
         # 使用增强版WebUI
        print("正在启动增强版 AI 搜索助手 Web 界面...")
        bot, use_es_retrieval = init_agent_service(retrieval_config)
        
        retrieval_method = "Elasticsearch 检索" if use_es_retrieval else "默认检索"
        print(f"当前使用的检索方式: {retrieval_method}")
        
        chatbot_config = {
            'prompt.suggestions': [
                '介绍下雇主责任险',
                '雇主责任险和工伤保险有什么主要区别？',
                '介绍一下平安商业综合责任保险的保障范围。',
                '施工保主要适用于哪些场景？',
                '请比较不同保险产品的优缺点',
                '请用数学公式解释保险费率的计算方法'
            ]
        }
    
        enhanced_ui = EnhancedWebUI(bot, chatbot_config=chatbot_config, use_es_retrieval=use_es_retrieval)
           
        enhanced_ui.run()
    except Exception as e:
        print(f"启动 Web 界面失败: {e}")
        print("请检查网络连接、API Key 以及 Elasticsearch 服务是否正常运行。")


def main_cli(doc_threshold: int = 50, force_es_retrieval: bool = False):
    """
    启动命令行交互界面
    :param doc_threshold: 文档数量阈值
    :param force_es_retrieval: 是否强制使用ES检索
    """
    try:
        # 创建检索配置
        es_config = {
            "host": "https://localhost",
            "port": 9200,
            "user": "elastic",
            "password": "ncUg27UP5aMSzSElVXAX",
            "index_name": "adaptive_search_docs_index"
        }
        
        retrieval_config = AdaptiveRetrievalConfig(
            doc_threshold=doc_threshold,
            use_es_retrieval=force_es_retrieval,
            es_config=es_config
        )
        
        print("正在启动 AI 搜索助手 CLI 界面...")
        bot, use_es_retrieval = init_agent_service(retrieval_config)
        
        retrieval_method = "Elasticsearch 检索" if use_es_retrieval else "默认检索"
        print(f"当前使用的检索方式: {retrieval_method}")
        print("输入 'quit' 或 'exit' 退出程序")
        
        # 对话历史
        messages = []
        while True:
            try:
                query = input('\n请输入您的问题: ').strip()
                if query.lower() in ['quit', 'exit', '退出']:
                    print("再见！")
                    break
                
                if not query:
                    print("问题不能为空！")
                    continue
                
                messages.append({'role': 'user', 'content': query})
                print("正在处理您的问题...")
                
                # 流式输出回答
                response = []
                for response_chunk in bot.run(messages=messages):
                    if response_chunk and response_chunk[-1]['role'] == 'assistant':
                        assistant_message = response_chunk[-1]
                        print(assistant_message.get('content', ''), end='', flush=True)
                    response = response_chunk
                
                print()  # 换行
                messages.extend(response)
                
            except KeyboardInterrupt:
                print("\n程序被用户中断")
                break
            except Exception as e:
                print(f"\n处理问题时出错: {e}")
                
    except Exception as e:
        print(f"启动 CLI 界面失败: {e}")


# 增强版WebUI类，用于在插件框下方添加知识库信息
class EnhancedWebUI(WebUI):
    """增强版 WebUI，在插件框下面添加知识库信息"""
    
    def __init__(self, agent, chatbot_config=None,use_es_retrieval=False, **kwargs):
        super().__init__(agent, chatbot_config, **kwargs)
        # 获取 docs 文件夹下的真实文件列表
        self.knowledge_bases = self._get_docs_files()
        self.use_es_retrieval = use_es_retrieval  # 是否使用 Elasticsearch 检索
        
    def _create_ui(self):
        """重写 _create_ui 方法，添加知识库信息"""
        ui = super()._create_ui()
        ui["plugins"].append(self._create_knowledge_base_block)
        return ui
    
    def _get_docs_files(self):
        """获取 docs 文件夹下的文件列表"""
        file_dir = os.path.join(os.path.dirname(__file__), 'docs')
        files = []
        if os.path.exists(file_dir):
            for file in os.listdir(file_dir):
                file_path = os.path.join(file_dir, file)
                if os.path.isfile(file_path):
                    files.append(file)
        return sorted(files)  # 按文件名排序
 
    def _create_knowledge_base_block(self):
        """创建知识库信息块"""
        from qwen_agent.gui.gradio_dep import gr
        
        # 不使用 Group，直接创建组件，避免背景色
        gr.Markdown("### 📚 保险知识库")
                
        # 显示前4个文件的完整文件名（无序号）
        gr.Markdown(f"**已加载{len(self.knowledge_bases)}个文档：**")
        display_count = min(4, len(self.knowledge_bases))
        for i in range(display_count):
            gr.Markdown(f"• {self.knowledge_bases[i]}")
        
        # 如果还有更多文件，显示剩余数量
        remaining_count = len(self.knowledge_bases) - display_count
        if remaining_count > 0:
            gr.Markdown(f"还有{remaining_count}个文件未显示")
        
        # 添加检索方式信息
        retrieval_method = "Elasticsearch 检索" if self.use_es_retrieval else "默认检索"
        gr.Markdown(f"**当前检索方式：** {retrieval_method}")
        
        return None
    
    def run(self, messages=None, share=False, server_name=None, server_port=None, concurrency_limit=10, enable_mention=False, **kwargs):
        """重写 run 方法，添加知识库显示"""
        self.run_kwargs = kwargs

        from qwen_agent.gui.gradio_dep import gr, mgr, ms
        from qwen_agent.gui.utils import convert_history_to_chatbot

        customTheme = gr.themes.Default(
            primary_hue=gr.themes.utils.colors.blue,
            radius_size=gr.themes.utils.sizes.radius_none,
        )

        with gr.Blocks(
                css=os.path.join(os.path.dirname(__file__), 'assets/appBot.css') if os.path.exists(os.path.join(os.path.dirname(__file__), 'assets/appBot.css')) else None,
                theme=customTheme,
        ) as demo:
            history = gr.State([])
            with ms.Application():
                with gr.Row(elem_classes='container'):
                    with gr.Column(scale=4):
                        chatbot = mgr.Chatbot(value=convert_history_to_chatbot(messages=messages),
                                              avatar_images=[
                                                  self.user_config,
                                                  self.agent_config_list,
                                              ],
                                              height=850,
                                              avatar_image_width=80,
                                              flushing=False,
                                              show_copy_button=True,
                                              ##配置laTex数学公式的渲染支持
                                              latex_delimiters=[{ 
                                                  'left': '\\(',
                                                  'right': '\\)',
                                                  'display': True
                                              }, {
                                                  'left': '\\begin{equation}',
                                                  'right': '\\end{equation}',
                                                  'display': True
                                              }, {
                                                  'left': '\\begin{align}',
                                                  'right': '\\end{align}',
                                                  'display': True
                                              }, {
                                                  'left': '\\begin{alignat}',
                                                  'right': '\\end{alignat}',
                                                  'display': True
                                              }, {
                                                  'left': '\\begin{gather}',
                                                  'right': '\\end{gather}',
                                                  'display': True
                                              }, {
                                                  'left': '\\begin{CD}',
                                                  'right': '\\end{CD}',
                                                  'display': True
                                              }, {
                                                  'left': '\\[',
                                                  'right': '\\]',
                                                  'display': True
                                              }])

                        input = mgr.MultimodalInput(placeholder=self.input_placeholder,)
                        audio_input = gr.Audio(
                            sources=["microphone"],
                            type="filepath"
                        )

                    with gr.Column(scale=1):
                        if len(self.agent_list) > 1:
                            agent_selector = gr.Dropdown(
                                [(agent.name, i) for i, agent in enumerate(self.agent_list)],
                                label='Agents',
                                info='选择一个Agent',
                                value=0,
                                interactive=True,
                            )

                        agent_info_block = self._create_agent_info_block()
                        agent_plugins_block = self._create_agent_plugins_block()
                        
                        # 添加知识库信息块
                        knowledge_base_block = self._create_knowledge_base_block()

                        if self.prompt_suggestions:
                            gr.Examples(
                                label='推荐对话',
                                examples=self.prompt_suggestions,
                                inputs=[input],
                            )

                    if len(self.agent_list) > 1:
                        agent_selector.change(
                            fn=self.change_agent,
                            inputs=[agent_selector],
                            outputs=[agent_selector, agent_info_block, agent_plugins_block],
                            queue=False,
                        )

                    input_promise = input.submit(
                        fn=self.add_text,
                        inputs=[input, audio_input, chatbot, history],
                        outputs=[input, audio_input, chatbot, history],
                        queue=False,
                    )

                    if len(self.agent_list) > 1 and enable_mention:
                        input_promise = input_promise.then(
                            self.add_mention,
                            [chatbot, agent_selector],
                            [chatbot, agent_selector],
                        ).then(
                            self.agent_run,
                            [chatbot, history, agent_selector],
                            [chatbot, history, agent_selector],
                        )
                    else:
                        input_promise = input_promise.then(
                            self.agent_run,
                            [chatbot, history],
                            [chatbot, history],
                        )

                    input_promise.then(self.flushed, None, [input])

            demo.load(None)

        demo.queue(default_concurrency_limit=concurrency_limit).launch(share=share,
                                                                       server_name=server_name,
                                                                       server_port=server_port)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='AI 搜索助手')
    parser.add_argument('--mode', choices=['gui', 'cli'], default='gui', 
                        help='运行模式: gui(图形界面) 或 cli(命令行界面)')
    parser.add_argument('--threshold', type=int, default=50,
                        help='文档数量阈值，超过此数量使用ES检索 (默认: 50)')
    parser.add_argument('--force-es', action='store_true',
                        help='强制使用 Elasticsearch 检索')
    
    args = parser.parse_args()
    
    if args.mode == 'gui':
       
        main(doc_threshold=args.threshold, force_es_retrieval=args.force_es)
        
    else:
        main_cli(doc_threshold=args.threshold, force_es_retrieval=args.force_es)